from p4.v1 import p4runtime_pb2
import grpc
from p4runtime import P4RuntimeClient

def setup_p4runtime(switches):
    for switch in switches:
        client = P4RuntimeClient(switch['addr'], switch['device_id'])
        client.master_arbitration()
        # 定义所有主机的转发规则（以 s1、s2、s3 为例，主机 h1-h3 跨交换机部署）
        tunnel_ids = {
            's1-s2': 100, 's2-s1': 101,
            's1-s3': 200, 's3-s1': 201,
            's2-s3': 300, 's3-s2': 301,
        }
        host_mappings = {
            'h1': {'switch': 's1', 'port': 1, 'ip': '10.0.0.1'},
            'h2': {'switch': 's2', 'port': 1, 'ip': '10.0.0.2'},
            'h3': {'switch': 's3', 'port': 1, 'ip': '10.0.0.3'},
        }
        # 为每个交换机配置流规则和计数器
        for src_host, src_info in host_mappings.items():
            for dst_host, dst_info in host_mappings.items():
                if src_host == dst_host:
                    continue
                # 确定隧道 ID 和转发方向
                tunnel_key = f"{src_info['switch']}-{dst_info['switch']}"
                reverse_tunnel_key = f"{dst_info['switch']}-{src_info['switch']}"
                tunnel_id = tunnel_ids[tunnel_key]
                reverse_tunnel_id = tunnel_ids[reverse_tunnel_key]
                
                # 步骤1：配置正向流规则（src→dst）
                table_entry = p4runtime_pb2.TableEntry(
                    table_name="MyIngress.forward",
                    action_name="MyIngress.set_tunnel",
                )
                # 匹配 IPv4 目的地址和隧道 ID
                table_entry.match.append(
                    p4runtime_pb2.FieldMatch(
                        field_id=client.get_field_id("hdr.ipv4.dstAddr"),
                        exact=p4runtime_pb2.Exact(value=dst_info['ip'].encode())
                    )
                )
                table_entry.match.append(
                    p4runtime_pb2.FieldMatch(
                        field_id=client.get_field_id("meta.tunnel_id"),
                        exact=p4runtime_pb2.Exact(value=tunnel_id.to_bytes(4, 'big'))
                    )
                )
                # 关联正向计数器
                table_entry.action.params.append(
                    p4runtime_pb2.Action.Param(
                        param_id=client.get_param_id("tunnel_counter_id"),
                        value=tunnel_id.to_bytes(4, 'big')
                    )
                )
                client.write_table_entry(table_entry)
                
                # 步骤2：配置反向流规则（dst→src）
                reverse_table_entry = p4runtime_pb2.TableEntry(
                    table_name="MyIngress.forward",
                    action_name="MyIngress.set_tunnel",
                )
                reverse_table_entry.match.append(
                    p4runtime_pb2.FieldMatch(
                        field_id=client.get_field_id("hdr.ipv4.dstAddr"),
                        exact=p4runtime_pb2.Exact(value=src_info['ip'].encode())
                    )
                )
                reverse_table_entry.match.append(
                    p4runtime_pb2.FieldMatch(
                        field_id=client.get_field_id("meta.tunnel_id"),
                        exact=p4runtime_pb2.Exact(value=reverse_tunnel_id.to_bytes(4, 'big'))
                    )
                )
                reverse_table_entry.action.params.append(
                    p4runtime_pb2.Action.Param(
                        param_id=client.get_param_id("tunnel_counter_id"),
                        value=reverse_tunnel_id.to_bytes(4, 'big')
                    )
                )
                client.write_table_entry(reverse_table_entry)

        client.close()

# 步骤2：采集计数器并写入文件
def collect_counters(switches, output_files):
    for switch, file in zip(switches, output_files):
        client = P4RuntimeClient(switch['addr'], switch['device_id'])
        client.master_arbitration()
        tunnel_ids = {100, 101, 200, 201, 300, 301}
        with open(file, 'w') as f:
            for tid in tunnel_ids:
                # 读取 ingress 计数器（接收方向）
                ingress_counter = client.read_counter("MyIngress.ingressTunnelCounter", tid)
                # 读取 egress 计数器（发送方向）
                egress_counter = client.read_counter("MyIngress.egressTunnelCounter", tid)
                f.write(f"隧道 {tid} 接收包数：{ingress_counter.packet_count}, 字节数：{ingress_counter.byte_count}\n")
                f.write(f"隧道 {tid} 发送包数：{egress_counter.packet_count}, 字节数：{egress_counter.byte_count}\n")
        client.close()

# 主函数执行
if __name__ == "__main__":
    switches = [
        {'addr': 'localhost:50051', 'device_id': 0, 'name': 's1'},
        {'addr': 'localhost:50052', 'device_id': 1, 'name': 's2'},
        {'addr': 'localhost:50053', 'device_id': 2, 'name': 's3'},
    ]
    output_files = ['S1S2.txt', 'S1S3.txt', 'S2S3.txt']
    setup_p4runtime(switches)
    collect_counters(switches, output_files)